import arrow

from pymongo import MongoClient, UpdateOne
from conf.settings import CLOUDMONITOR_MONGODB, CLOUDMONITOR_DB

conn = MongoClient(CLOUDMONITOR_MONGODB)
db = conn[CLOUDMONITOR_DB]

t0 = "2021-11-01T00:00:00+08:00"
t1 = "2022-03-01T00:00:00+08:00"


def period_statistics():
    task_type_tuple = ('securityEvent', 'vul')

    query = {'endTime': {'$gt': arrow.get(t0).datetime, '$lt': arrow.get(t0).shift(months=1).datetime}, 'taskType': {'$in': task_type_tuple}}

    job_id_list = list(set([item.get('jobId') for item in db.tasks.find(query)]))

    query_1 = {
        'jobId': {'$in': job_id_list},
        'endTime': {'$gt': arrow.get(t1).datetime, '$lt': arrow.get(t1).shift(months=1).datetime},
        'taskType': {'$in': task_type_tuple}
    }

    job_id_set = set([item.get('jobId') for item in db.tasks.find(query_1)])

    job_id_list_2 = list(job_id_set)

    print(len(job_id_list_2))

    query['jobId'] = {'$in': job_id_list_2}
    query_1['jobId'] = {'$in': job_id_list_2}

    for _query in (query, query_1):
        print('===============')
        s_count = 0
        vul_count = 0
        user_list = []
        d0 = {}
        i = 0
        for item in db.tasks.find(_query):
            user_list.append(item.get('uid'))

            task_type = item.get('taskType')

            if task_type == 'vul':
                vul_count += 1
                d1 = get_vul(item)
                for k, v in d1.items():
                    d0.setdefault(k, 0)
                    d0[k] += d1[k]
            else:
                s_count += 1
                i += get_securityEvent(item)

        print(f"用户数: {len(set(user_list))}")
        print(f"资产数: {len(job_id_list_2)}")
        print(f"漏洞：{d0}")
        print(f"安全事件数: {i}")
        print(f"扫描安全事件次数: {s_count}")
        print(f"扫描漏洞次数：{vul_count}")


def get_vul(item):
    d0 = {}

    for warning in item.get('result', {}).get('warnings', []):
        lv_name = warning.get('levelName')
        count = d0.get(lv_name, 0)
        d0[lv_name] = count + 1

    return d0


def get_securityEvent(item):
    return len(item.get('result', {}).get('warnings', []))


def statistics_2():
    low = "2022-03-01T00:00:00+08:00"
    high = "2022-04-01T00:00:00+08:00"

    uid_list = [item.get('uid') for item in db.balances.find({'planExpireTime': {'$gt': arrow.get(low).datetime}})]

    task_type_dict = {}
    uid_dict = {}
    job_dict = {}
    for item in db.tasks.find(
            {'endTime': {'$gte': arrow.get(low).datetime, '$lt': arrow.get(high).datetime}, 'uid': {'$in': uid_list}},
            {'_id': 0, 'taskType': 1, 'uid': 1, 'jobId': 1}
    ):
        task_type = item.get('taskType')
        uid = item.get('uid')
        job_id = item.get('jobId')

        task_type_dict[task_type] = task_type_dict.get(task_type, 0) + 1
        uid_dict[uid] = 1
        job_dict[job_id] = 1

    print(task_type_dict)
    print("job_id", len(job_dict))
    print("uid", len(uid_dict))


def main():
    for ti in (t0, t1):
        d0 = {}
        for item in db.tasks.find({'endTime': {'$gt': arrow.get(ti).datetime, '$lt': arrow.get(ti).shift(months=1).datetime}, 'taskType': 'vul'}):
            for warning in item.get('result', {}).get('warnings', []):
                l_n = warning.get('levelName')
                x = d0.get(l_n, 0)
                d0[l_n] = x + 1

        print(ti)
        print(d0)


if __name__ == '__main__':
    period_statistics()
